package weiyao.xinxidasai.Utils.Encryption.Chaos;

import cn.hutool.crypto.CryptoException;
import org.apache.commons.codec.digest.DigestUtils;
import weiyao.xinxidasai.DTO.VerificationReport;
import weiyao.xinxidasai.pojo.Version;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Arrays;
import java.util.List;

public class NotaCryptor {
    private static final String HMAC_ALGORITHM = "HmacSHA256";
    private static final int SALT_LENGTH = 16;
    private static final int HMAC_LENGTH = 32;
    private static final int TOKEN_HEADER_LENGTH = 32;
    private static final int ORIGINAL_LENGTH_HEADER = 4;
    private static final int FILE_ID_LENGTH = 8; // 文件ID长度（8字节）

    // 生成带隐写token的头部
    private static byte[] generateHeaderWithToken(String token, int originalLength) throws Exception {
        if (token == null || token.isEmpty()) {
            throw new IllegalArgumentException("Token cannot be null or empty");
        }

        // 1. 生成盐值
        byte[] salt = new byte[SALT_LENGTH];
        SecureRandom.getInstanceStrong().nextBytes(salt);

        // 2. 加密token
        byte[] tokenBytes = token.getBytes(StandardCharsets.UTF_8);
        byte[] encryptedToken = encryptToken(tokenBytes, salt);

        // 3. 组合头部 [盐值 + 加密后的token + 原始长度]
        ByteArrayOutputStream headerStream = new ByteArrayOutputStream();
        headerStream.write(salt);
        headerStream.write(encryptedToken);

        // 添加原始长度信息 (4字节，大端序)
        headerStream.write((byte) (originalLength >>> 24));
        headerStream.write((byte) (originalLength >>> 16));
        headerStream.write((byte) (originalLength >>> 8));
        headerStream.write((byte) originalLength);

        return headerStream.toByteArray();
    }

    // 使用盐值加密token
    private static byte[] encryptToken(byte[] token, byte[] salt) {
        byte[] encrypted = new byte[TOKEN_HEADER_LENGTH];
        for (int i = 0; i < TOKEN_HEADER_LENGTH; i++) {
            if (i < token.length) {
                encrypted[i] = (byte) (token[i] ^ salt[i % salt.length]);
            } else {
                encrypted[i] = salt[i % salt.length];
            }
        }
        return encrypted;
    }

    // 增强的加密方法
    public static void encryptWithToken(InputStream inputStream, OutputStream outputStream,
                                        Version version, String token) throws CryptoException {
        try {
            // 读取原始数据
            ByteArrayOutputStream dataBuffer = new ByteArrayOutputStream();
            byte[] buffer = new byte[1024];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                dataBuffer.write(buffer, 0, bytesRead);
            }
            byte[] originalData = dataBuffer.toByteArray();
            int originalLength = originalData.length;

            // 生成带token的头部
            byte[] header = generateHeaderWithToken(token, originalLength);
            outputStream.write(header);

            // 生成混沌密钥流
            CEOLogisticMap.CEOMap ceoMap = new CEOLogisticMap.CEOMap(
                    version.getR(), version.getX0(), version.getIterations());
            List<Byte> keyStream = ceoMap.generateKeyStream();

            // 异或加密
            ByteArrayOutputStream xorOutput = new ByteArrayOutputStream();
            for (int i = 0; i < originalData.length; i++) {
                byte keyByte = keyStream.get(i % keyStream.size());
                xorOutput.write(originalData[i] ^ keyByte);
            }
            byte[] encryptedData = xorOutput.toByteArray();

            // 写入加密数据
            outputStream.write(encryptedData);

            // 计算并写入HMAC
            ByteArrayOutputStream fullData = new ByteArrayOutputStream();
            fullData.write(header);
            fullData.write(encryptedData);
            byte[] hmac = calculateHMAC(fullData.toByteArray(), token);
            outputStream.write(hmac);

        } catch (Exception e) {
            throw new CryptoException("加密失败: " + e.getMessage(), e);
        }
    }

    // 增强的解密方法
    public static VerificationReport decryptWithTokenRecovery(InputStream inputStream,
                                                              Version version,
                                                              String expectedToken) throws CryptoException {
        VerificationReport report = new VerificationReport();
        try {
            // [1] 读取头部 (52字节)
            int headerLength = SALT_LENGTH + TOKEN_HEADER_LENGTH + ORIGINAL_LENGTH_HEADER;
            byte[] header = new byte[headerLength];
            int headerBytes = readFully(inputStream, header, 0, headerLength);
            if (headerBytes != headerLength) {
                throw new CryptoException("头部不完整");
            }

            // [2] 解析头部信息
            byte[] salt = Arrays.copyOfRange(header, 0, SALT_LENGTH);
            byte[] encryptedToken = Arrays.copyOfRange(header, SALT_LENGTH, SALT_LENGTH + TOKEN_HEADER_LENGTH);

            // 解析原始长度
            int originalLength = ((header[SALT_LENGTH + TOKEN_HEADER_LENGTH] & 0xFF) << 24) |
                    ((header[SALT_LENGTH + TOKEN_HEADER_LENGTH + 1] & 0xFF) << 16) |
                    ((header[SALT_LENGTH + TOKEN_HEADER_LENGTH + 2] & 0xFF) << 8) |
                    (header[SALT_LENGTH + TOKEN_HEADER_LENGTH + 3] & 0xFF);
            report.setOriginalLength(originalLength);

            // 解密token
            byte[] decrypted = new byte[TOKEN_HEADER_LENGTH];
            for (int i = 0; i < TOKEN_HEADER_LENGTH; i++) {
                decrypted[i] = (byte) (encryptedToken[i] ^ salt[i % salt.length]);
            }
            int length = 0;
            while (length < decrypted.length && decrypted[length] != 0) {
                length++;
            }
            String embeddedToken = new String(decrypted, 0, length, StandardCharsets.UTF_8);
            report.setEmbeddedToken(embeddedToken);

            // [3] 读取加密数据 (尝试读取原始长度，但允许不完整)
            byte[] encryptedData = new byte[originalLength];
            int dataRead = readFully(inputStream, encryptedData, 0, originalLength);

            // 记录实际读取的字节数
            report.setRecoveredBytes(dataRead);

            // [4] 尝试读取HMAC (32字节)，但允许失败
            byte[] storedHmac = new byte[HMAC_LENGTH];
            int hmacRead = readFully(inputStream, storedHmac, 0, HMAC_LENGTH);
            boolean hmacComplete = (hmacRead == HMAC_LENGTH);

            // 计算并验证HMAC (如果可能)
            boolean hmacValid = false;
            if (hmacComplete) {
                ByteArrayOutputStream fullData = new ByteArrayOutputStream();
                fullData.write(header);
                fullData.write(encryptedData, 0, dataRead);
                byte[] calculatedHmac = calculateHMAC(fullData.toByteArray(), expectedToken);
                hmacValid = Arrays.equals(storedHmac, calculatedHmac);
            }

            // [5] 解密数据 (只解密实际读取的部分)
            CEOLogisticMap.CEOMap ceoMap = new CEOLogisticMap.CEOMap(
                    version.getR(), version.getX0(), version.getIterations());
            List<Byte> keyStream = ceoMap.generateKeyStream();

            ByteArrayOutputStream decryptedOutput = new ByteArrayOutputStream();
            for (int i = 0; i < dataRead; i++) {
                // 循环使用密钥流
                byte keyByte = keyStream.get(i % keyStream.size());
                decryptedOutput.write(encryptedData[i] ^ keyByte);
            }
            byte[] decryptedData = decryptedOutput.toByteArray();
            report.setDecryptedData(decryptedData);

            // [6] 验证报告
            boolean tokenValid = expectedToken.equals(embeddedToken);
            String message = tokenValid ? "Token验证通过" : "Token不匹配: " + embeddedToken;

            if (!hmacComplete) {
                message = "HMAC不完整! " + message;
            } else if (!hmacValid) {
                message = "HMAC验证失败! " + message;
            }

            int qualityScore = calculateQualityScore(hmacValid, tokenValid,
                    decryptedData.length, originalLength);

            report.setValid(tokenValid && hmacValid && dataRead == originalLength);
            report.setTampered(!tokenValid || !hmacValid || dataRead != originalLength);
            report.setQualityScore(qualityScore);
            report.setMessage(message);
            report.setRecoveredBytes(dataRead);

            return report;

        } catch (Exception e) {
            report.setValid(false);
            report.setTampered(true);
            report.setQualityScore(0);
            report.setMessage("解密失败: " + e.getMessage());
            report.setRecoveredBytes(0);
            report.setDecryptedData(new byte[0]);
            return report;
        }
    }

    // 安全读取方法（允许读取不完整数据）
    private static int readFully(InputStream in, byte[] buffer, int offset, int length) throws IOException {
        int totalRead = 0;
        while (totalRead < length) {
            int bytesRead = in.read(buffer, offset + totalRead, length - totalRead);
            if (bytesRead == -1) {
                break; // 流结束
            }
            totalRead += bytesRead;
        }
        return totalRead;
    }

    // 计算质量评分 (0-100)
    private static int calculateQualityScore(boolean hmacValid, boolean tokenValid,
                                             int recoveredBytes, int originalLength) {
        int score = 0;
        if (tokenValid) score += 50;
        if (hmacValid) score += 25;

        if (originalLength > 0 && recoveredBytes == originalLength) {
            score += 25;
        } else if (originalLength > 0 && recoveredBytes > 0) {
            double recoveryRate = (double) recoveredBytes / originalLength;
            score += (int)(25 * recoveryRate);
        }
        return Math.min(100, score);
    }

    private static byte[] calculateHMAC(byte[] data, String token) throws NoSuchAlgorithmException, InvalidKeyException {
        Mac mac = Mac.getInstance(HMAC_ALGORITHM);
        javax.crypto.spec.SecretKeySpec secretKey = new SecretKeySpec(token.getBytes(StandardCharsets.UTF_8), HMAC_ALGORITHM);
        mac.init(secretKey);
        return mac.doFinal(data);
    }
}